from itertools import chain
from .base import *

FILE_CS = "cspk"
CS_ATTRIB = ["cnt","c","d","v2g"]

class StaCSPK(StaBase):
    def __init__(self,path:str,suffix:str,tinst:TrafficInst,plugins:dict[str,PluginBase]):
        head = cross_list(tinst.get_sta_head(),["cnt","c","pb"]) + cross_list(tinst.PKList.get_CS_names(),["ps","d","v2g"])
        super().__init__(FILE_CS,path,suffix,head,tinst,plugins)

    def GetData(self,inst:TrafficInst,plugins:dict[str,PluginBase])->list:
        t = inst.current_time
        cnt = [float(f) for f in inst.get_veh_count()]
        Pc = [cs.Pc_kW for cs in chain(inst.CSList,inst.PKList)]
        pb = [cs.pbuy(t) for cs in chain(inst.CSList,inst.PKList)]
        if inst.PKList[0].supports_V2G:
            ps = [cs.psell(t) for cs in inst.PKList]
        else:
            ps = [0] * len(inst.PKList)
        Pd = [cs.Pd_kW for cs in inst.PKList]
        Pv2g = [cs.Pv2g_kW for cs in inst.PKList]
        return cnt + Pc + pb + ps + Pd + Pv2g